package com.example.grpc.client.rpc;

import com.example.grpc.client.config.GrpcClientProperties;
import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;
import jakarta.annotation.PreDestroy;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

@Component
public class GrpcChannelManager {
    private static final Logger logger = LoggerFactory.getLogger(GrpcChannelManager.class);
    private final Map<String, ManagedChannel> channels = new ConcurrentHashMap<>();
    private final GrpcClientProperties properties;
    
    public GrpcChannelManager(GrpcClientProperties properties) {
        this.properties = properties;
    }
    
    public ManagedChannel getChannel(String serviceName) {
        return channels.computeIfAbsent(serviceName, this::createChannel);
    }
    
    public ManagedChannel getChannel(Endpoint endpoint) {
        String key = endpoint.getServiceName() + ":" + endpoint.getHost() + ":" + endpoint.getPort();
        return channels.computeIfAbsent(key, k -> createChannel(endpoint));
    }
    
    private ManagedChannel createChannel(String serviceName) {
        GrpcClientProperties.ServiceConfig config = properties.getServices().get(serviceName);
        if (config == null) {
            throw new IllegalArgumentException("No configuration found for service: " + serviceName);
        }
        
        return createChannelBuilder(config.getHost(), config.getPort(), config)
                .build();
    }
    
    private ManagedChannel createChannel(Endpoint endpoint) {
        GrpcClientProperties.ServiceConfig config = properties.getServices()
                .getOrDefault(endpoint.getServiceName(), new GrpcClientProperties.ServiceConfig());
        
        // 使用endpoint的配置，如果properties中没有配置
        String host = config.getHost() != null ? config.getHost() : endpoint.getHost();
        Integer port = config.getPort() != null ? config.getPort() : endpoint.getPort();
        
        return createChannelBuilder(host, port, config)
                .build();
    }
    
    private ManagedChannelBuilder<?> createChannelBuilder(String host, Integer port,
                                                          GrpcClientProperties.ServiceConfig config) {
        ManagedChannelBuilder<?> builder = ManagedChannelBuilder.forAddress(host, port);
        
        if (config.isPlaintext()) {
            builder.usePlaintext();
        }
        
        return builder
                .keepAliveTime(config.getKeepAliveTimeNanos(), TimeUnit.NANOSECONDS)
                .keepAliveTimeout(config.getKeepAliveTimeoutNanos(), TimeUnit.NANOSECONDS)
                .keepAliveWithoutCalls(config.isKeepAliveWithoutCalls())
                .maxInboundMessageSize(config.getMaxInboundMessageSize());
    }
    
    @PreDestroy
    public void shutdown() {
        logger.info("Shutting down gRPC channels...");
        channels.values().forEach(channel -> {
            try {
                channel.shutdown().awaitTermination(5, TimeUnit.SECONDS);
            } catch (InterruptedException e) {
                logger.warn("Failed to shutdown channel gracefully", e);
                channel.shutdownNow();
                Thread.currentThread().interrupt();
            }
        });
        channels.clear();
    }
}